from apyori import apriori


class BookApriori(object):
    @staticmethod
    def list_flatten(data):
        """
        将二维数组转一维数组，并将其转为不可变set
        :param data:
        :return:
        """
        c = [i for j in data for i in j]
        c_set = list(set(c))
        c_set.sort()
        # 每个元素即为一个不可变集合
        return [frozenset([num]) for num in c_set]
        # return map(frozenset, [])

    @staticmethod
    def get_freq_item_set(data, candidate_set, min_support):
        """
        获取满足条件(支持度 不小于 最小支持度)的频繁项集

        :param data: 数据集
        :param candidate_set: 频繁项集的候选项集
        :param min_support:  最小支持度
        :return:
        """
        candidate_set_dict = {}
        for info in data:
            for candidate in candidate_set:
                if candidate.issubset(set(info)):
                    candidate_set_dict[candidate] = candidate_set_dict.get(candidate, 0) + 1

        support_dict = {}
        freq_item = []
        for k, v in candidate_set_dict.items():
            support_value = v / len(data)
            if support_value >= min_support:
                freq_item.append(k)
                support_dict[frozenset(k)] = support_value
        return freq_item, support_dict

    @staticmethod
    def compute_freq_item(freq_item_set, k):
        """

        :param freq_item_set:
        :param k:
        :return:
        """
        freq_item = []
        for num, item in enumerate(freq_item_set):
            for i in range(num + 1, len(freq_item_set)):
                # 前k-2个项相同时，将两个集合合并，生成新的候选集合
                l1 = list(item)[: k-2]
                l2 = list(freq_item_set[i])[: k-2]
                l1.sort()
                l2.sort()
                if l1 == l2:
                    # 并集
                    freq_item.append(item|freq_item_set[i])
        return freq_item

    def apriori(self, data, min_support=0.5):
        """

        :param data:
        :param min_support:
        :return:
        """
        # 二维数组转一维数组，并转为不可变集合
        candidate = self.list_flatten(data)
        freq_item, support_data = self.get_freq_item_set(data, candidate_set=candidate, min_support=min_support)

        freq_item_total = [freq_item]
        k = 2
        # 当频繁项集不为空
        while len(freq_item_total[k-2]) > 0:
            candidate_new = self.compute_freq_item(freq_item_total[k-2], k)
            mid_freq_item, mid_support_data = self.get_freq_item_set(data, candidate_set=candidate_new, min_support=min_support)
            # 更新频繁项集及其支持度
            support_data.update(mid_support_data)
            freq_item_total.append(mid_freq_item)
            k += 1
        return freq_item_total, support_data

    @staticmethod
    def calc_credible(freq_set, freq_items, support_data, br1, min_credible=0.7):
        """
        计算可信度，返回可信度大于阈值的所有频繁项集元素的集合（对于频繁项集(a, b),可以推出a -- > b， 而不支持b --> a）

        :param freq_set: 频繁项集
        :param freq_items: 繁项集子集
        :param support_data: 支持度dict
        :param br1: 头两个为并集的集合，第三个值为两个集合并集后的可信度值
        :param min_credible: 最小可信度
        :return:
        """
        pruned_freq_items = []
        for item in freq_items:
            # 可信度： p和h并集的支持度 / p的支持度
            credible = support_data[frozenset(freq_set)] / support_data[frozenset(set(freq_set) - item)]
            if credible >= min_credible:
                print("{} --> {}, credible: {}".format(freq_set - item, item, credible))
                # 如果符合条件，则在频繁项集freq_set中，子集freq_set - item --> item 的推理是可信的（大概率）
                br1.append((freq_set - item, item, credible))
                pruned_freq_items.append(item)
        return pruned_freq_items

    def rules_from_freq_item(self, freq_set, freq_items, support_data, br1, min_credible=0.7):
        """
        生成满足条件(>=最小可信度)的所有频繁项集的子集

        :param freq_set: 频繁项集
        :param freq_items: 繁项集子集
        :param support_data: 支持度dict
        :param br1:
        :param min_credible: 最小可信度
        :return:
        """
        if len(freq_set) > (len(freq_items[0]) + 1):
            # 从频繁项集中，按指定长度生成所有的频繁项集子集
            freq_items_new = self.compute_freq_item(freq_items, len(freq_items[0]) + 1)
            # 按最低可信度过滤满足条件的频繁项集子集
            freq_items_new = self.calc_credible(freq_set, freq_items_new, support_data, br1, min_credible)
            # 若满足条件的频繁项集子集依然存在，则再次循环（频繁项集子集长度+1，直到循环完成所有的子集）
            if len(freq_items_new) > 1:
                self.rules_from_freq_item(freq_set, freq_items_new, support_data, br1, min_credible)

    def generate_association_rule(self, freq_set_list, support_data, min_credible=0.7):
        """
        生成所有符合条件的关联规则，包括置信度

        :param freq_set_list:  频繁项集列表
        :param support_data: 支持度dict
        :param min_credible: 最小可信度
        :return:
        """
        rule_list = []
        for num, freq_set in enumerate(freq_set_list):
            if num == 0:
                # 单元素子集，无法推理
                continue
            for freq_info in freq_set:
                # 逐项循环生成频繁项集子集,计算子集间推理的可信度是否超过阈值
                freq_items = [frozenset([item]) for item in freq_info]
                self.calc_credible(freq_info, freq_items, support_data, br1=rule_list, min_credible=min_credible)
                # 当子集元素大于2个时，将符合条件（置信度判断）的子集合并生成新的子集
                if num > 1:
                    self.rules_from_freq_item(freq_info, freq_items, support_data, br1=rule_list, min_credible=min_credible)

                # if num > 1:
                #     self.rules_from_freq_item(freq_info, freq_items, support_data, br1=rule_list, min_credible=min_credible)
                # else:
                #     self.calc_credible(freq_info, freq_items, support_data, br1=rule_list, min_credible=min_credible)
        return rule_list


class MyApriori(object):
    """
    获取数据集的频繁项集 Apriori算法
    """

    def __init__(self, data, min_support):
        self.data = data
        self.min_support = min_support

    @staticmethod
    def is_item_in_line(line, item):
        """
        判断item(单元素或数组)是否是line的子集
        :param line:
        :param info:
        :return:
        """
        for info in item:
            if isinstance(info, frozenset):
                if list(info)[0] not in line:
                    return False
            else:
                if info not in line:
                    return False
        return True

    def apriori_code(self, candidata_freq):
        """
        单元素的频繁项集
        :param candidata_freq: 频繁项集候选集
        :return:
        """
        support_item = dict()
        for line in self.data:
            for item in candidata_freq:
                if not isinstance(item, list):
                    item = [item]
                if self.is_item_in_line(line, item):
                    support_item[frozenset(item)] = support_item.get(frozenset(item), 0) + \
                                                    1 / len(self.data)

        mid_support = support_item.copy()
        for k, v in mid_support.items():
            if v < self.min_support:
                del support_item[k]

        return support_item

    @staticmethod
    def get_random_list(arr):
        """
        数组中每两个子集合并，以生成获取新的子集维度+1维的数组
        :param arr:
        :param num:
        :return:
        """
        if len(arr) < 2:
            return arr

        new_arr = []
        for i in range(len(arr) - 1):
            for j in range(i + 1, len(arr)):
                mid_arr = list(set(arr[i]) | set(arr[j]))
                mid_arr.sort()
                if len(mid_arr) == len(arr[0]) + 1 and mid_arr not in new_arr:
                    new_arr.append(mid_arr)

        return new_arr

    def apriori_code_main(self, condidata_freq, freq_items):
        """
        apriori主程序

        :param condidata_freq: 候选频繁项集
        :param freq_items: 频繁项集结果集
        :return:
        """
        if not condidata_freq:
            candidata_items = [i for j in self.data for i in j]
            new_candidata_items = list(set(candidata_items))

        if isinstance(condidata_freq, dict):
            new_candidata_items = list(condidata_freq.keys())
            new_candidata_items = self.get_random_list(new_candidata_items)

        mid_freq_items = self.apriori_code(candidata_freq=new_candidata_items)

        if not mid_freq_items:
            return freq_items

        freq_items.update(mid_freq_items)
        return self.apriori_code_main(condidata_freq=mid_freq_items, freq_items=freq_items)

    @staticmethod
    def calc_credible(freq_set, freq_items, support_data, br1, min_credible=0.7):
        """
        计算可信度，返回可信度大于阈值的所有频繁项集元素的集合（对于频繁项集(a, b),可以推出a -- > b， 而不支持b --> a）

        :param freq_set: 频繁项集
        :param freq_items: 繁项集子集
        :param support_data: 支持度dict
        :param br1: 头两个为并集的集合，第三个值为两个集合并集后的可信度值
        :param min_credible: 最小可信度
        :return:
        """
        pruned_freq_items = []
        for item in freq_items:
            # 可信度： p和h并集的支持度 / p的支持度
            credible = support_data[frozenset(freq_set)] / support_data[frozenset(set(freq_set) - set(item))]
            if credible >= min_credible:
                print("{} --> {}, credible: {}".format(freq_set - set(item), item, credible))
                # 如果符合条件，则在频繁项集freq_set中，子集freq_set - item --> item 的推理是可信的（大概率）
                br1.append((freq_set - set(item), item, credible))
                pruned_freq_items.append(item)
        return pruned_freq_items

    def rules_from_freq_item(self, freq_set, freq_items, support_data, br1, min_credible=0.7):
        """
        生成满足条件(>=最小可信度)的所有频繁项集的子集

        :param freq_set: 频繁项集
        :param freq_items: 繁项集子集
        :param support_data: 支持度dict
        :param br1:
        :param min_credible: 最小可信度
        :return:
        """
        if len(freq_set) > (len(freq_items[0]) + 1):
            # 从频繁项集中，按指定长度生成所有的频繁项集子集
            # freq_items_new = self.compute_freq_item(freq_items, len(freq_items[0]) + 1)
            freq_items_new = self.get_random_list(arr=freq_items)
            # 按最低可信度过滤满足条件的频繁项集子集
            freq_items_new = self.calc_credible(freq_set, freq_items_new, support_data, br1, min_credible)
            # 若满足条件的频繁项集子集依然存在，则再次循环（频繁项集子集长度+1，直到循环完成所有的子集）
            if len(freq_items_new) > 1:
                self.rules_from_freq_item(freq_set, freq_items_new, support_data, br1, min_credible)

    def generate_association_rule(self, freq_set_list, support_data, min_credible=0.7):
        """
        生成所有符合条件的关联规则，包括置信度

        :param freq_set_list:  频繁项集列表
        :param support_data: 支持度dict
        :param min_credible: 最小可信度
        :return:
        """
        rule_list = []
        for num, freq_set in enumerate(freq_set_list):
            if num == 0:
                # 单元素子集，无法推理
                continue
            for freq_info in freq_set:
                # 逐项循环生成频繁项集子集,计算子集间推理的可信度是否超过阈值
                freq_items = [frozenset([item]) for item in freq_info]
                self.calc_credible(freq_info, freq_items, support_data, br1=rule_list,
                                   min_credible=min_credible)
                # 当子集元素大于2个时，将符合条件（置信度判断）的子集合并生成新的子集
                if num > 1:
                    self.rules_from_freq_item(freq_info, freq_items, support_data, br1=rule_list,
                                              min_credible=min_credible)

        return rule_list


def load_data(file):
    with open(file, 'r') as f:
        file_data = f.readlines()

    # arr = [info[: -1].split('\t') for info in file_data]
    arr = [info[: -1].strip().split(' ') for info in file_data]
    # 国会选举
    # dict_info = dict()
    # for info in file_data:
    #     data = info[: -1].split('\t')
    #     dict_info[int(data[0])] = data[1]
    # return dict_info
    return [list(map(int, info)) for info in arr if info]


def run_simple():
    book_apriori = BookApriori()

    data = [[1, 3, 4], [2, 3, 5], [1, 2, 3, 5], [2, 5]]
    freq_item, support_data = book_apriori.apriori(data, min_support=0.5)
    print('freq_item: {}'.format(freq_item))
    print('support_data: {}'.format(support_data))
    rules = book_apriori.generate_association_rule(freq_set_list=freq_item, support_data=support_data, min_credible=0.5)
    print('rules: {}'.format(rules))


def run_mushroom():
    # file_recent100bills = './data/B/chp11/recent100bills.txt'
    # file_recent20bills = './data/B/chp11/recent20bills.txt'
    #
    # recent100bills = load_data(file_recent100bills)
    # file_recent20bills = load_data(file_recent20bills)
    #
    # print(recent100bills)
    # print(file_recent20bills)
    # 第一个特征：1：可食用；2：有毒
    # 第二个特征：蘑菇伞的形状 3-8来标识
    mushroom_file = './data/chp11/mushroom.dat'
    mushroom_data = load_data(mushroom_file)

    book_apriori = BookApriori()
    freq_item, support_data = book_apriori.apriori(mushroom_data, min_support=0.3)

    for item in freq_item[1]:
        if 2 in item:
            print(item)


def run_compare_3():
    data = [['I1', 'I2', 'I5'],
            ['I2', 'I4'],
            ['I2', 'I3'],
            ['I1', 'I2', 'I4'],
            ['I1', 'I3'],
            ['I2', 'I3'],
            ['I1', 'I3'],
            ['I1', 'I2', 'I3', 'I5'],
            ['I1', 'I2', 'I3']]

    min_support_ratio = 2 / len(data)
    min_confidence = 0.7

    # my code
    print('my code')
    my_apriori = MyApriori(data=data, min_support=min_support_ratio)

    freq_items = my_apriori.apriori_code_main(condidata_freq=None, freq_items=dict())
    print('my_apriori freq_items: {}'.format(freq_items))
    print(len(freq_items))
    freq_set = []
    for i in range(1, len(set([i for j in data for i in j])) + 1):
        mid_arr = []
        for k in list(freq_items.keys()):
            if len(k) == i:
                mid_arr.append(k)
        if mid_arr:
            freq_set.append(mid_arr)
    # print('freq_set: {}'.format(freq_set))
    my_rule_list = my_apriori.generate_association_rule(freq_set_list=freq_set, support_data=freq_items,
                                                        min_credible=min_confidence)
    print('my_rule_list:{}'.format(my_rule_list))
    print(len(my_rule_list))

    # apyori包
    print('\napyori包 代码')
    res = list(apriori(transactions=data, min_support=min_support_ratio, min_confidence=min_confidence))

    for info in res:
        print(info)
    print(len(res))

    # book code
    print('\ncode from book')
    book_apriori = BookApriori()
    freq_item, support_data = book_apriori.apriori(data, min_support=min_support_ratio)
    print('freq_item: {}'.format(freq_item))
    print('support_data: {}'.format(support_data))
    rules = book_apriori.generate_association_rule(freq_set_list=freq_item, support_data=support_data,
                                                   min_credible=min_confidence)
    print('rules: {}'.format(rules))
    print(len(rules))


if __name__ == '__main__':
    # run_simple()
    # run_mushroom()
    run_compare_3()
